import cv2
import os
import numpy
from scipy.sparse import data
import torch
import pickle 
import config as cfg
from tqdm import tqdm
from os import listdir, name, rename
from os.path import join
from numpy import array, zeros, save, rot90
from models.facenet.facenet import Facenet
from face_recognition import face_locations, load_image_file
from sklearn.neighbors import KNeighborsClassifier


def encoding(img_path, del_bad_pic=True):
    paths = listdir(img_path)
    features = []
    names = []
    warn_msgs = []

    # 加载模型
    path = 'models/facenet/facenet_mobilenet.pth'
    model = Facenet(mode='predict').eval()
    state_dict = torch.load(path, map_location='cpu')
    model.load_state_dict(state_dict, strict=False)

    # 遍历人脸图片
    for i, path in zip(tqdm(range(len(paths))), paths):
        image = load_image_file(join(img_path, path))
        
        image = cv2.resize(image, (720, 1280))
        # iPhone拍照要转九十度, 其他情况请注释该行代码
        image = rot90(image, 3)
        faces = []
        boxes = face_locations(image, model='hog')
        if len(boxes) != 1:
            # 人脸不等于1， 存入消息表，运行完成后统一输出。
            warn_msgs.append("file {0} can't found face, or found more than 1 face".format(join(img_path, path)))
            if del_bad_pic:
                os.remove(join(img_path, path))
            continue
        else:
            # 将人脸特征向量与人名写入列表。
            box = boxes[0]
            face = image[box[0]:box[2], box[3]:box[1]]
            face = cv2.cvtColor(face, cv2.COLOR_RGB2BGR)
            face = face / 255.
            faces.append(cv2.resize(face, (160, 160)))
            faces = array(faces).transpose(0, 3, 1, 2)
            faces = torch.from_numpy(faces).type(torch.FloatTensor)
            feature = model(faces).detach().numpy()
            features.append(feature[0])
            name = path.split('_')[0]
            names.append(name)
    
    if del_bad_pic:
        format_imgname(img_path)
        
    names = array(names)

    return features, names, warn_msgs

def format_imgname(img_path):
    paths = listdir(img_path)
    name_table = {}
    # 排序
    for p in paths:
        name = p.split('_')[0]
        if name in name_table.keys():
            name_table[name].append(p)
        else:
            name_table[name] = [p]
    # 重命名
    for img_list in name_table.keys():
        i = 1
        img_list = name_table[img_list]
        img_list = sorted(img_list, key=lambda x: int(x.split('_')[1].split('.')[0]))
        for img in img_list:
            rename(join(img_path, img), join(img_path, "{0}_{1}.{2}".format(img.split('_')[0], i, img.split('.')[1])))
            i += 1        


if __name__ == '__main__':
    print('start encoding faces')
    features, names, warn_msgs = encoding(cfg.pictures_of_konw)
    numpy.save('data/names.npy', names)
    numpy.save('data/features.npy', features)
    neigh = KNeighborsClassifier(n_neighbors=3, algorithm='ball_tree', weights='distance')
    neigh.fit(features, names)
    with open('data/knnd3.pickle', 'wb') as f:
        pickle.dump(neigh, f)
    print('encode faces complete, take {} warnings:'.format(len(warn_msgs)))
    for msg in warn_msgs:
        print('\twarnings:{}'.format(msg))